import torch
from util import extract, linear_beta_schedule
import matplotlib.pyplot as plt
from load_model import load_model

def sample(model, Config):
    #设置随机种子
    torch.manual_seed(Config.seed)
    betas = linear_beta_schedule(Config.timesteps)
    alphas = 1. - betas
    alphas_cumprod = torch.cumprod(alphas, axis=0)  #前缀积
    sqrt_recip_alphas = torch.sqrt(1. / alphas)
    sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - alphas_cumprod)
    model.eval()
    with torch.no_grad():
        # Generate sample from noise
        sample = torch.randn(1, Config.channels, Config.image_size, Config.image_size).to(Config.device)
        for t in reversed(range(0, Config.timesteps)):
            t_batch = torch.full((1,), t, device=Config.device, dtype=torch.long)
            predicted_noise = model(sample, t_batch)
            
            # Update sample according to reverse process
            beta_t = extract(betas, t_batch, sample.shape)
            sqrt_one_minus_alphas_cumprod_t = extract(sqrt_one_minus_alphas_cumprod, t_batch, sample.shape)
            sqrt_recip_alphas_t = extract(sqrt_recip_alphas, t_batch, sample.shape)
            
            sample = sqrt_recip_alphas_t * (
                sample - beta_t * predicted_noise / sqrt_one_minus_alphas_cumprod_t
            )
            noise = torch.randn_like(sample)
            sqrt_beta_t = extract(torch.sqrt(betas), t_batch, sample.shape)
            sample += sqrt_beta_t * noise
        
        # Plot the generated sample
        sample = sample.squeeze(0).permute(1, 2, 0).cpu().numpy()
        sample = (sample * 0.5 + 0.5).clip(0, 1)  # Scale from [-1,1] to [0,1]
        
        plt.figure(figsize=(4,4))
        plt.imshow(sample)
        plt.title(f"Sample with seed {Config.seed}")
        plt.axis('off')
        plt.savefig(f"sample_seed_{Config.seed}.png")
        plt.close()

if __name__=="__main__":
    from Config import Config
    model=load_model("chosen_model/32pic_model.pth")
    sample(model,Config)